use ndarray::{Array1, Array2};

#[allow(unused)]
pub enum IterationMethod {
    Jacobi,
    GaussSeidel,
    SOR(f64), // 松弛因子 ω
}

impl IterationMethod {
    pub fn solve(&self, x: Array2<f64>, y: Array1<f64>) -> Array1<f64> {
        let n = x.shape()[0];
        let mut x_new = Array1::<f64>::zeros(n);
        let mut x_old = Array1::<f64>::zeros(n);

        let max_iter = 1000;
        let tol = 1e-6;

        for _ in 0..max_iter {
            for i in 0..n {
                let mut sigma = 0.0;
                for j in 0..n {
                    if j != i {
                        sigma += x[[i, j]] * match self {
                            IterationMethod::Jacobi => x_old[j],
                            IterationMethod::GaussSeidel | IterationMethod::SOR(_) => x_new[j],
                        };
                    }
                }

                let new_value = (y[i] - sigma) / x[[i, i]];
                x_new[i] = match self {
                    IterationMethod::Jacobi => new_value,
                    IterationMethod::GaussSeidel => new_value,
                    IterationMethod::SOR(omega) => {
                        (1.0 - omega) * x_old[i] + omega * new_value
                    }
                };
            }

            if (&x_new - &x_old).mapv(|v| v.abs()).sum() < tol {
                break;
            }

            x_old.assign(&x_new);
        }

        x_new
    }
}


#[cfg(test)]
#[allow(unused)]
mod test {
    use super::*;
    use ndarray::array;
    use approx::assert_abs_diff_eq;

    #[test]
    fn test_solve() {
        let a = array![
            [4.0, 1.0, 2.0],
            [3.0, 5.0, 1.0],
            [1.0, 1.0, 3.0]
        ];
        let b = array![4.0, 7.0, 3.0];
        
        let ans = array![0.5, 1.0, 0.5];
        let epsilon= 1e-6;

        let x = IterationMethod::Jacobi.solve(a.clone(), b.clone());
        assert_abs_diff_eq!(
            x.as_slice().unwrap(), 
            ans.as_slice().unwrap(),
            epsilon = epsilon
        );

        let x = IterationMethod::GaussSeidel.solve(a.clone(), b.clone());
                assert_abs_diff_eq!(
            x.as_slice().unwrap(), 
            ans.as_slice().unwrap(),
            epsilon = epsilon
        );

        let x = IterationMethod::SOR(1.5).solve(a, b);
        assert_abs_diff_eq!(
            x.as_slice().unwrap(), 
            ans.as_slice().unwrap(),
            epsilon = epsilon
        );
    }
}
